# 本文件用于训练 "训练集调用与训练" 一步
import struct
import torch


fd = open("./MNIST/MNIST/train-images.idx3-ubyte", "rb")

# ->0123
print(fd.read(4)) # read 0 to 3 byte

# 从文件第4字节开始读取
fd.seek(4) # 0123 ->4567

# 将文件中的值以byte形式读取然后转化为int类型识别
# >i 指 大端-整形数据
# z = struct.unpack('>i', fd.read(4))[0]
num = struct.unpack('>i', fd.read(4))[0]
x = struct.unpack('>i', fd.read(4))[0]
y = struct.unpack('>i', fd.read(4))[0]

print(num, x, y) # read 4 to 7 byte

# 创建一个数组存放图片数据
arr = torch.zeros(x, y)
for imgx in range(x):
    for imgy in range(y):
        arr[imgx, imgy] = struct.unpack('B', fd.read(1))[0]
print(arr)

########################################
# 一、nn.Module（网络的构建）
#
########################################
# 导入包
import torch.nn as nn
import torch.nn.functional as F

# 创类
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()   # 初始化子类 （能否理解为构造函数）
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)
    # 前向传播
    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))


########################################
# 二、torch.nn.functional.conv2d（卷积运算）
#
########################################
# 设置输入矩阵
inputs = torch.tensor([[1, 2, 0, 3, 1],
                       [0, 1, 2, 3, 1],
                       [1, 2, 1, 0, 0],
                       [5, 2, 3, 1, 1],
                       [2, 1, 0, 1, 1]])
# 设置卷积核大小
# kernel = torch.zeros(12, 12)
kernel = torch.tensor([[1, 2, 1],
                       [0, 1, 0],
                       [2, 1, 0]])

#转换成函数参数要求的形状
inputs = torch.reshape(inputs, (1, 1, 5, 5))
kernel = torch.reshape(kernel, (1, 1, 3, 3))

# 调用函数
output = F.conv2d(inputs, kernel, stride=1)
print(output)

########################################
# 三、神经网络–卷积层（一层卷积网络的构建）
#
########################################

import torchvision
from torch.utils.data import DataLoader

# 从官网下载数据
# dataset = torchvision.datasets.CIFAR10('./dataset/CIFAR10', transform=torchvision.transforms.ToTensor(), download=True)
dataset = torchvision.datasets.MNIST("./MNIST/", transform=torchvision.transforms.ToTensor())

# 加载数据
dataloader = DataLoader(dataset, batch_size=64)


# 实例化网络
model = Model()

#对数据进行卷积
for data in dataloader:
    img, target = data
    output = model(img)
    print(output.shape)

torch.save(model.state_dict(), "./save/a.pth")

